
#include "../Base.h"
typedef NTSTATUS(*FileCallBack)(HANDLE FileHandle, PVOID info, PVOID un1, PVOID un2);
typedef struct _RegisterCallback {
	FileCallBack QueryFileCallback;
	FileCallBack SetFileCallBack;
}RegisterCallback, * PRegisterCallback;


typedef NTSTATUS(*ExRegisterAttrbuteInformationCallbackProc)(PRegisterCallback callback);



FileCallBack oldExpDisQueryAttributeInformation = NULL;
FileCallBack oldExpDisSetAttributeInformation = NULL;


NTSTATUS QueryFileCallBack(HANDLE FileHandle, PVOID info, PVOID un1, PVOID un2) {

	LOG("xxxxxxx\r\n");
	if (MmIsAddressValid(info)) {
		PCommPackage package = (PCommPackage)info;
		if (package->id == 0x12345678) {
			package->retStatus = gCommCallback(package);
		}
		else {

			if (oldExpDisQueryAttributeInformation) {
				return oldExpDisQueryAttributeInformation(FileHandle, info, un1, un2);
			}
		}
	}

	return STATUS_SUCCESS;
}

NTSTATUS SetFileCallBack(HANDLE FileHandle, PVOID info, PVOID un1, PVOID un2) {
	LOG("xxxxxxx\r\n");

	if (MmIsAddressValid(info)) {
		PCommPackage package = (PCommPackage)info;
		if (package->id == 0x12345678) {
			package->retStatus = gCommCallback(package);
		}
		else {

			if (oldExpDisSetAttributeInformation) {
				return oldExpDisSetAttributeInformation(FileHandle, info, un1, un2);
			}
		}
	}

	return STATUS_SUCCESS;
}


VOID UnRegisterCommWin7() {

	gWin7CommCallBack[0] = (ULONG64)oldExpDisQueryAttributeInformation;
	gWin7CommCallBack[1] = (ULONG64)oldExpDisSetAttributeInformation;
}



NTSTATUS RegisterCommWin7(CommCallbackProc callback) {
	UNICODE_STRING uFuncName = { 0 };
	RtlInitUnicodeString(&uFuncName, L"ExRegisterAttributeInfomationCallback");
	PUCHAR pFunc = (PUCHAR)MmGetSystemRoutineAddress(&uFuncName);
	ULONG64 offset = *(PLONG)(pFunc + 0xd + 3);
	PULONG64 ExpDisQueryAttributeInformation = (PULONG64)((pFunc + 0xd + 7) + offset);
	oldExpDisQueryAttributeInformation = (FileCallBack)ExpDisQueryAttributeInformation[0];
	oldExpDisSetAttributeInformation = (FileCallBack)ExpDisQueryAttributeInformation[1];
	ExpDisQueryAttributeInformation[0] = 0;
	ExpDisQueryAttributeInformation[1] = 0;

	RegisterCallback rcb = { 0 };
	rcb.QueryFileCallback = QueryFileCallBack;
	rcb.SetFileCallBack = SetFileCallBack;
	ExRegisterAttrbuteInformationCallbackProc callFunc = (ExRegisterAttrbuteInformationCallbackProc)pFunc;
	NTSTATUS status = callFunc(&rcb);
	if (!NT_SUCCESS(status)) {
		return STATUS_UNSUCCESSFUL;
	}
	gCommCallback = callback;
	gWin7CommCallBack = ExpDisQueryAttributeInformation;
	return STATUS_SUCCESS;
}